1904E - Tree Queries - CodeForces Solution


binary search data structures dfs and similar implementation trees

Please click on ads to support us..

C++ Code:

#include <bits/stdc++.h>
using namespace std;
#define mod 1000000007
#define oo 1000000010
 
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
#define getrand(l, r) uniform_int_distribution<int>(l, r)(rng)
 
const int N = 200010;
const int LOG = 19;
 
 
int n , q;
 
vector< int > g[N];
 
int st[LOG][N] , depth[N] , dfs_l[N] , dfs_r[N] , dfs_cnt = 0 , at[N];
 
void DFS(int node,int prnt,int depth){
	dfs_l[node] = dfs_r[node] = dfs_cnt++;
	::depth[node] = depth;
	st[0][node] = prnt;
	for(int k = 1;k < LOG;k++){
		if(st[k - 1][node] == -1)
			st[k][node] = -1;
		else
			st[k][node] = st[k - 1][st[k - 1][node]];
	}
	for(int i = 0 ;i < (int)g[node].size();i++){
		if(g[node][i] == prnt)
			continue;
		DFS(g[node][i] , node , depth + 1);
		dfs_r[node] = dfs_r[g[node][i]];
	}
	at[dfs_l[node]] = node;
}
 
long long seg[4 * N] , lazy[4 * N];
 
inline void fix(int s,int e,int idx){
	seg[idx] += lazy[idx];
	if(s != e){
		lazy[(idx << 1)] += lazy[idx];
		lazy[(idx << 1) + 1] += lazy[idx];
	}
	lazy[idx] = 0;
}
 
long long build(int s,int e,int idx){
	if(s == e){
		return seg[idx] = depth[at[s]];
	}
	return seg[idx] = max(build(s, ((s+e) >> 1),(idx << 1)),build(((s+e) >> 1) + 1,e,(idx << 1) + 1));
}
 
long long update(int s,int e,int idx,int l,int r,int val){
	fix(s , e, idx);
	if(s > r || e < l)
		return seg[idx];
	if(s >= l && e <= r){
		lazy[idx] += val;
		fix(s , e , idx);
		return seg[idx];
	}
	return seg[idx] = max(update(s, ((s+e) >> 1),(idx << 1), l , r , val),update(((s+e) >> 1) + 1,e,(idx << 1) + 1 ,l  ,r , val));
}
 
long long getmax(int s,int e,int idx,int l,int r){
	fix(s , e, idx);
	if(s > r || e < l)
		return -oo;
	if(s >= l && e <= r)
		return seg[idx];
	return max(getmax(s, ((s+e) >> 1),(idx << 1), l , r),getmax(((s+e) >> 1) + 1,e,(idx << 1) + 1 ,l  ,r));
}
 
int LCA(int u,int v){
	if(depth[u] > depth[v])
		swap(u , v);
	for(int k = LOG - 1;k >= 0;k--){
		if(depth[v] - (1 << k) >= depth[u])
			v = st[k][v];
	}
	if(u == v)
		return u;
	for(int k = LOG - 1;k >= 0;k--){
		if(st[k][u] != st[k][v]){
			u = st[k][u];
			v = st[k][v];
		}
	}
	return st[0][u];
}
 
int arr[N] , seg2[4 * N];
 
int update2(int s,int e,int idx,int Idx,int val){
	if(s > Idx || e < Idx)
		return seg2[idx];
	if(s == e)
		return seg2[idx] = val;
	return seg2[idx] = max(update2(s, ((s+e) >> 1),(idx << 1), Idx , val) , update2(((s+e) >> 1) + 1,e,(idx << 1) + 1,  Idx , val));
}
 
int getmax2(int s,int e,int idx,int l,int r){
	if(s > r || e < l)
		return -oo;
	if(s >= l && e <= r)
		return seg2[idx];
	return max(getmax2(s , ((s+e) >> 1) , (idx << 1) , l , r),getmax2(((s+e) >> 1) + 1,e,(idx << 1) + 1, l , r));
}
 
vector< pair< int , vector< int > > > qu[N];
 
int ans[N];
 
int bl[N] , br[N];
 
void calc_node(int node){
	long long tmp = 0;
	if(bl[node] == -1)
		tmp = getmax(0 , n - 1 , 1 , dfs_l[node] , dfs_r[node]);
	else{
		tmp = getmax(0 , n - 1 , 1 , dfs_l[node] , bl[node] - 1);
		if(br[node] + 1 <= dfs_r[node])
			tmp = max(tmp , getmax(0 , n - 1 , 1 , br[node] + 1 , dfs_r[node]));
	}
	//cout << "calc node " << node << " " << depth[node] << " " << tmp << " " << bl[node] << " " << br[node] << " " << (int)tmp - 2 * depth[node] << endl;
	update2(0 , n - 1 , 1 , depth[node] , (int)tmp - 2 * depth[node]);
}
 
void DFS(int node,int prnt){
	bl[node] = br[node] = -1;
	calc_node(node);
	for(int i = 0 ;i < (int)g[node].size();i++){
		if(g[node][i] == prnt)
			continue;
		bl[node] = dfs_l[g[node][i]];
		br[node] = dfs_r[g[node][i]];
		calc_node(node);
		DFS(g[node][i] , node);
	}
	bl[node] = br[node] = -1;
	calc_node(node);
 
	for(int mxd = 0 , i = 0 ;i < (int)qu[node].size();i++){
		//cout << "answer " << node << endl;
		vector< int > &tmp = qu[node][i].second;
		mxd = 0;
		for(int j = 0 ;j < (int)tmp.size();j++){
			if(dfs_l[tmp[j]] <= dfs_l[node] && dfs_r[tmp[j]] >= dfs_l[node]){
				mxd = max(mxd , depth[tmp[j]] + 1);
			}
			else{
				//cout << "remove node " << tmp[j] << " with lca " << LCA(node , tmp[j]) << endl;
				update(0 , n - 1 , 1 , dfs_l[tmp[j]] , dfs_r[tmp[j]] , -oo);
				calc_node(LCA(node , tmp[j]));
			}
		}
		//cout << "mxd = " << mxd << " " << depth[node] << " " << getmax2(0 , n - 1 , 1 , mxd , depth[node]) << endl;
		ans[qu[node][i].first] = getmax2(0 , n - 1 , 1 , mxd , depth[node]) + depth[node];
		for(int j = 0 ;j < (int)tmp.size();j++){
			if(!(dfs_l[tmp[j]] <= dfs_l[node] && dfs_r[tmp[j]] >= dfs_l[node])){
				update(0 , n - 1 , 1 , dfs_l[tmp[j]] , dfs_r[tmp[j]] , oo);
				calc_node(LCA(node , tmp[j]));
			}
		}
	}
}
 
int main(){
	scanf("%d%d",&n,&q);
	for(int a , b , i = 0 ;i < n - 1;i++){
		scanf("%d%d",&a,&b);
		g[a].push_back(b);
		g[b].push_back(a);
	}
	DFS(1 , -1 , 0);
	build(0 , n - 1 , 1);
	int x , k;
	for(int i = 0 ;i < q;i++){
		scanf("%d%d",&x,&k);
		qu[x].push_back(make_pair(i , vector< int > (k)));
		for(int j = 0 ;j < k;j++)
			scanf("%d",&qu[x].back().second[j]);
	}
	DFS(1 , -1);
	for(int i = 0 ;i < q;i++)
		printf("%d\n",ans[i]);
	return 0;
}


Comments

Submit
0 Comments
More Questions

866. Prime Palindrome
1516A - Tit for Tat
622. Design Circular Queue
814. Binary Tree Pruning
791. Custom Sort String
787. Cheapest Flights Within K Stops
779. K-th Symbol in Grammar
701. Insert into a Binary Search Tree
429. N-ary Tree Level Order Traversal
739. Daily Temperatures
647. Palindromic Substrings
583. Delete Operation for Two Strings
518. Coin Change 2
516. Longest Palindromic Subsequence
468. Validate IP Address
450. Delete Node in a BST
445. Add Two Numbers II
442. Find All Duplicates in an Array
437. Path Sum III
436. Find Right Interval
435. Non-overlapping Intervals
406. Queue Reconstruction by Height
380. Insert Delete GetRandom O(1)
332. Reconstruct Itinerary
368. Largest Divisible Subset
377. Combination Sum IV
322. Coin Change
307. Range Sum Query - Mutable
287. Find the Duplicate Number
279. Perfect Squares